If this is going to be the base class for all constraints I'd favor naming it Constraint

This is just used in the tests so I'd suggest moving it into the testing file and making it a private method. Unwrapping coordinates is quite tricky to do right, and I don't want to imply we've done it perfectly by adding it to the public API.

for now I can do that, but I consider that this implement should be right, at least assuming that the displacement at each step is small enough

ik this is a helper function, but if we're using it to verify the correctness of other code, it'd be nice to have a test for this

For 2. AtomIndexedConstraints spanning multiple systems:
So actually the current FixAtoms is implemented so that it is supposed to be only one FixAtoms contraints for a batch system. So it's expected to act on multiple systems.

for 1. it's not so clear that it's a problem to have the same atoms affected by multiple constraints. Another precise example of that, take a water molecule, you often want to constraints H-bonds. Then the oxygen atom will be concerned by 2 constraints.

Can we do this in a single comprehension with an if statement and walrus operator, like below?

since we want people to use this new function, maybe we should add a warning if they don't use set_positions?

@property
def positions(self) -> torch.Tensor:
    """Unit cell following the row vector convention."""
    logger.warn("Directly setting state.positions is deprecated. use set_positions instead")
    return self.positions

Or we can make the set_positions function the @Property of def positions instead so it will automatically handle constraints and people don't ever need to know about the set_positions function. We may need to prefix positions with an underscore since it the variable might shadow the @property definition but I'm not sure

I think this is a small but important thing to do because we don't want people to accidentally use the wrong functions to set positions/forces.

Ahh yeah this is important. I agree that having set_positions and positions is a bit confusing.

I would suggest we instead call this constrain_positions which has no arguments and applies the constraints to the existing positions. I don't think setting positions should necessarily incur the application of constraints, and having the naming mixed up is confusing. A separate operation strikes me as the cleanest solution. We could call it apply_positions constraints too.

Whatever we do, let's make sure to propagate it to forces and momenta as well.

So the flow would look like:

state.positions = new_positions
state.constrain_positions()

Actually this set_positions mostly makes sense for under-the-hood algorithms. I don't think you can add @Property on something that is a already a variable of the class, otherwise you just replace it. So your approach would work with a hidden _positions but this start to change a lot the current logic. By the way, you don't always want to apply constraints.
Example: in MD, it's redundant to apply on both forces and momenta. It's probably not wrong, though it may not be this simple for more complicated Constraint like FixBonds, but you do two time the work.

Ok upon reflection my proposed solution would not work. I suggest instead:

1. Removing the set_positions, set_momenta, and set_forces functions from the states
2. Instead creating constrain_positions, constrain_momenta, and constrain_forces functions in constraints.py, these would have the following signature and use:

def constrain_positions(positions: torch.Tensor, constraints: list[Constraint], reference_state: ts.SimState) -> torch.Tensor:
        for constraint in (constraints or []):
            constraint.adjust_positions(reference_state, positions)

# then used by
constrain_positions(new_positions, state.constraints, state)
state.positions = new_positions

# of course with a slight adjustment we could also do
constrain_positions(new_positions, state)
state.positions = new_positions

I think removing the constrain_X logic from the SimStates is preferable because it makes it more reusable across the package and isolates the contraint logic in constraints.py.

Potentially I would prefer a more verbose
state.set_positions(new_positions, apply_constraints=True). Only reason is having 1 line when you want to set a new position.
But I don't have a strong opinion on any of the different implementations we discussed

Yeah I agree having a one liner is nice. What about:

from constraints import set_constrained_positions

set_constrained_positions(state, new_positions)

I prefer not having two ways to do the same thing and I like having all the constraints logic in constraints. Could call it something else, I'd be ok with simply set_positions too.

I've changed my mind. One liner like ASE is a good option when you have a private variable _positions and force final user to use set_positions. In our case, it's not possible. So I think it's better as you said to keep in 2 lines

state.positions = positions
state.constrain_positions()

The code is more verbose but more self explainable as well.
I think I prefer this form as one understand that constraints come from state and modify internal positions, rather than constrain_positions which I think is a bit weird either with the argument/signature constrain_positions(new_positions, state.constraints, state) or constrain_positions(new_positions, state). In one case, it's a bit weird to provide state (but it's needed for some constraints (eg FixAtoms). The second case, the signature is the same as state.constraint_positions().

Sounds good to me! I could also imagine there could be special states that need their own constraint implementations, so I am happy to leave it on the state. It's a pretty short function, anyhow.

not sure what this mask= keyword is. but maybe this is leftover from older code. I think we can end this warning at "Perhaps you want to specify a mask instead"

I actually added support for mask. These init functions were copied from ASE a bit too quickly.

Ahh yeah this is important. I agree that having set_positions and positions is a bit confusing.

I would suggest we instead call this constrain_positions which has no arguments and applies the constraints to the existing positions. I don't think setting positions should necessarily incur the application of constraints, and having the naming mixed up is confusing. A separate operation strikes me as the cleanest solution. We could call it apply_positions constraints too.

Whatever we do, let's make sure to propagate it to forces and momenta as well.

So the flow would look like:

state.positions = new_positions
state.constrain_positions()